package cn.doitedu.ml.linear

import org.apache.spark.ml.evaluation.RegressionEvaluator
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.regression.LinearRegression
import org.apache.spark.sql.SparkSession

import scala.collection.mutable

/**
 * 线性回归代码示例
 */
object LinearDemo {

  def main(args: Array[String]): Unit = {


    val spark = SparkSession.builder()
      .appName("")
      .master("local")
      .getOrCreate()

    // 加载训练集
    val sample = spark.read.option("header", "true").csv("userprofile/data/linear/sample/sample.csv")

    val arr2Vec = (arr:mutable.WrappedArray[String])=>{
      Vectors.dense(arr.map(s=>s.toDouble).toArray)
    }
    spark.udf.register("arr2vec",arr2Vec)

    val sampleVec = sample.selectExpr("arr2vec(array(area,floor)) as features","cast(price as double) as price")

    // 训练模型
    val regression = new LinearRegression()
      .setFeaturesCol("features")
      .setLabelCol("price")

    val model = regression.fit(sampleVec)


    // 加载测试集
    val test = spark.read.option("header", "true").csv("userprofile/data/linear/test/test.csv")
    val testVec = test.selectExpr("arr2vec(array(area,floor)) as features","cast(price as double) as price")

    val result = model.transform(testVec)
    result.cache()
    result.show(100,false)

    val evaluator4 = new RegressionEvaluator()
      .setLabelCol("price")
      .setPredictionCol("prediction")
      .setMetricName("rmse")  //rmse  均方根误差   mse 均方误差  r2  mae：平均绝对误差
    val d4: Double = evaluator4.evaluate(result)
    println(d4)


    test.rdd.take(10)

    spark.close()
  }

}
